#! /usr/bin/env python
#coding=utf-8

import os
import sys

from dylib_file import DylibFile
sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), "../"))
from macho import DependencyBase

class Dependency(DependencyBase):
    def __init__(self, caller, callee):
        super(Dependency, self).__init__(caller, callee)

    def updateCallsFromDb(self, cursor):
        if self["calls"] > 0:
            return

        sqlcmd = "select count(symbol_id) from calls where dependence_id=%d" % (self["id"])
        cursor.execute(sqlcmd)
        for row in cursor:
            self["calls"] = row[0]
            return

    def queryCalls(self, cursor):
        querycmd = "select id from symbols where parent_id=%d and bits!='l' and name in (select name from undefines where parent_id=%d)" % (
        self["callee_id"], self["caller_id"])
        cursor.execute(querycmd)
        self["call_symbols"] = []
        for row in cursor:
            self["call_symbols"].append(row[0])
        self["calls"] = len(self["call_symbols"])

    def addMatchedCallsToDb(self, cursor):
        needCommit = False
        script = "BEGIN TRANSACTION;\n"
        for sym in self["call_symbols"]:
            sqlcmd = 'insert into calls(caller_id, callee_id, symbol_id, dependence_id) values (%d,%d,%d,%d)' % \
                    (self["caller_id"], self["callee_id"], sym, self["id"])
            script = script + sqlcmd + ";\n"
            needCommit = True
            if len(script) > 1024 * 4:
                script = script + "COMMIT;\n"
                cursor.executescript(script)
                needCommit = False
                script = "BEGIN TRANSACTION;\n"
        if needCommit:
            script = script + "COMMIT;\n"
            cursor.executescript(script)

        print("Add %d calls for %s->%s." % (self["calls"], self["caller"]["name"], self["callee"]["name"]))

if __name__ == '__main__':
    from macho import MachOFileMgr
    from dependency import Dependency
    mgr = MachOFileMgr(fileClass = DylibFile, dependenceClass=Dependency)
    mgr.scan_all_files()
    dylib = mgr.get_file_by_path("/System/Library/PrivateFrameworks/AppleFirmwareUpdate.framework/Versions/A/AppleFirmwareUpdate")

    import sqlite3
    _conn = sqlite3.connect(os.path.join(mgr.get_root_path(), "archinfo.db"))
    _cursor = _conn.cursor()

    dylib.addCalls(_cursor)

    _conn.commit()
    _cursor.close()
    _conn.close()
